package com.stark.commons.lang.util;

import com.stark.commons.lang.enums.MediaType;
import com.stark.commons.lang.util.HttpUtil.ByteArrayMultipartFile;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.codec.digest.DigestUtils;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.pdfbox.cos.COSObject;
import org.apache.pdfbox.pdmodel.DefaultResourceCache;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.pdmodel.encryption.AccessPermission;
import org.apache.pdfbox.pdmodel.encryption.PublicKeyRecipient;
import org.apache.pdfbox.pdmodel.encryption.StandardProtectionPolicy;
import org.apache.pdfbox.pdmodel.graphics.PDXObject;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.apache.pdfbox.text.PDFTextStripper;
import org.apache.pdfbox.tools.imageio.ImageIOUtil;

import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * Java 处理 PDF 方法集合。
 *
 * @author Ben
 * @since 1.0.0
 * @version 1.0.0
 * @see <a href="https://pdfbox.apache.org/">pdfbox官网</a>
 */
@Slf4j
public class PdfUtil {

	private static final float IMAGE_DPI = 300;

	/**
	 * 检查文件是否是 PDF 。
	 * @param bytes pdf 文件。
	 * @return 是返回 {@literal true} ，否返回 {@literal false} 。
	 */
	public static boolean isPdf(byte[] bytes) {
		if (ArrayUtils.isEmpty(bytes)) {
			return false;
		}

		PDDocument document = null;
		try {
			document = PDDocument.load(bytes);
			return true;
		} catch (IOException e) {
			return false;
		} finally {
			if (document != null) {
				try {
					document.close();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}
		}
	}

	/**
	 * 设置只读，不允许编辑或复制，允许打印。
	 * @param bytes pdf 文件。
	 * @return 处理后的 pdf 文件。
	 * @throws IOException
	 */
	public static byte[] setReadonly(byte[] bytes) throws IOException {
		PDDocument document = null;
		try {
			AccessPermission accessPermission = new AccessPermission();
			accessPermission.setCanAssembleDocument(false);
			accessPermission.setCanExtractContent(false);
			accessPermission.setCanExtractForAccessibility(false);
			accessPermission.setCanFillInForm(false);
			accessPermission.setCanModify(false);
			accessPermission.setCanModifyAnnotations(false);
			accessPermission.setCanPrint(true);
			accessPermission.setCanPrintDegraded(true);

			PublicKeyRecipient recipient = new PublicKeyRecipient();
			recipient.setPermission(accessPermission);

			StandardProtectionPolicy policy = new StandardProtectionPolicy(null, null, accessPermission);
			policy.setEncryptionKeyLength(128);

			document = PDDocument.load(bytes);
			document.protect(policy);
			ByteArrayOutputStream out = new ByteArrayOutputStream();
			document.save(out);
			bytes = out.toByteArray();
		} finally {
			if (document != null) {
				try {
					document.close();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}
		}
		return bytes;
	}

	/**
	 * 获取文本内容。
	 * @param bytes pdf 文件。
	 * @return 文本内容。
	 * @throws IOException 读取 pdf 出错时抛出异常。
	 */
	public static String getText(byte[] bytes) throws IOException {
		String text;
		PDDocument document = null;
		try {
			document = PDDocument.load(bytes);
			PDFTextStripper pdfStripper = new PDFTextStripper();
			text = pdfStripper.getText(document);
		} finally {
			if (document != null) {
				try {
					document.close();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}
		}
		return text;
	}

	/**
	 * pdf 转图片，每页一张图片。
	 * @param bytes pdf 文件。
	 * @param imageType 图片类型。
	 * @return 图片列表。
	 * @throws Exception 读取 pdf 出错时抛出异常。
	 */
	public static List<ByteArrayMultipartFile> pdfToImages(byte[] bytes, MediaType imageType) throws Exception {
		return pdfToImages(bytes, imageType, 0, 0);
	}

	/**
	 * pdf 转图片，每页一张图片。
	 * @param bytes pdf 文件。
	 * @param imageType 图片类型。
	 * @param dpi 像素。
	 * @return 图片列表。
	 * @throws Exception 读取 pdf 出错时抛出异常。
	 */
	public static List<ByteArrayMultipartFile> pdfToImages(byte[] bytes, MediaType imageType, float dpi) throws Exception {
		return pdfToImages(bytes, imageType, dpi, 0);
	}
	
	/**
	 * pdf 转图片，每页一张图片。
	 * @param bytes pdf 文件。
	 * @param imageType 图片类型。
	 * @param dpi 像素。
	 * @param maxSize 单张图片最大体积，单位字节，小于等于 0 表示不限制。
	 * @return 图片列表。
	 * @throws Exception 读取 pdf 出错时抛出异常。
	 */
	public static List<ByteArrayMultipartFile> pdfToImages(byte[] bytes, MediaType imageType, float dpi, long maxSize) throws Exception {
		if (!imageType.name().startsWith("IMAGE_")) {
			throw new IllegalAccessException("不支持的格式: " + imageType);
		}
		if (dpi <= 0) {
			dpi = IMAGE_DPI;
		}
		
		String id = DigestUtils.md5Hex(bytes);
		List<ByteArrayMultipartFile> images = new ArrayList<>();
		PDDocument document = null;
		try {
			document = PDDocument.load(bytes);
			int pages = document.getNumberOfPages();
			int length = String.valueOf(pages).length();
			PDFRenderer renderer = new PDFRenderer(document);
			BufferedImage image;
			ByteArrayOutputStream out;
			long START, start, end, consume;
			for (int i = 0; i < pages; i++) {
				START = start = System.currentTimeMillis();
				image = renderer.renderImageWithDPI(i, dpi);
				out = new ByteArrayOutputStream();
				ImageIOUtil.writeImage(image, imageType.getExtensions()[0], out);
				image.getGraphics().dispose();
				bytes = out.toByteArray();
				end = System.currentTimeMillis();
				consume = end - start;
				if (maxSize > 0) {
					start = System.currentTimeMillis();
					bytes = ImageUtil.autoSize(bytes, maxSize);
					end = System.currentTimeMillis();
				}
				log.debug("\n-->[" + id + "-" + pages + "-" + (i + 1) + "] PDF 转 " + imageType.getExtensions()[0].toUpperCase()
						+ "\n   转换图片耗时 " + consume + " 毫秒"
						+ (maxSize > 0 ? "\n   缩放耗时 " + (end - start) + " 毫秒" : "")
						+ "\n   总耗时 " + (end - START) + "毫秒");
				images.add(new ByteArrayMultipartFile("page_" + prependZero(i + 1, length) + "." + imageType.getExtensions()[0], bytes));
			}
		} finally {
			if (document != null) {
				try {
					document.close();
				} catch (IOException e) {
					// do nothing
				}
			}
		}
		return images;
	}

	/**
	 * pdf 转图片，每页一张图片。
	 * @param bytes pdf 文件。
	 * @param imageType 图片类型。
	 * @param dir 存储目录。
	 * @return 图片绝对路径列表。
	 * @throws Exception 读取 pdf 出错时抛出异常。
	 */
	public static List<String> pdfToImages(byte[] bytes, MediaType imageType, String dir) throws Exception {
		return pdfToImages(bytes, imageType, 0, 0, dir);
	}

	/**
	 * pdf 转图片，每页一张图片。
	 * @param bytes pdf 文件。
	 * @param imageType 图片类型。
	 * @param dpi 像素。
	 * @param dir 存储目录。
	 * @return 图片绝对路径列表。
	 * @throws Exception 读取 pdf 出错时抛出异常。
	 */
	public static List<String> pdfToImages(byte[] bytes, MediaType imageType, float dpi, String dir) throws Exception {
		return pdfToImages(bytes, imageType, dpi, 0, dir);
	}
	
	/**
	 * pdf 转图片，每页一张图片。
	 * @param bytes pdf 文件。
	 * @param imageType 图片类型。
	 * @param dpi 像素。
	 * @param maxSize 单张图片最大体积，单位字节，小于等于 0 表示不限制。
	 * @param dir 存储目录。
	 * @return 图片绝对路径列表。
	 * @throws Exception 读取 pdf 出错时抛出异常。
	 */
	public static List<String> pdfToImages(byte[] bytes, MediaType imageType, float dpi, long maxSize, String dir) throws Exception {
		if (!imageType.name().startsWith("IMAGE_")) {
			throw new IllegalAccessException("不支持的格式: " + imageType);
		}
		if (dpi <= 0) {
			dpi = IMAGE_DPI;
		}
		
		String id = DigestUtils.md5Hex(bytes);
		List<String> images = new ArrayList<>();
		PDDocument document = null;
		try {
			document = PDDocument.load(bytes);
			int pages = document.getNumberOfPages();
			int length = String.valueOf(pages).length();
			PDFRenderer renderer = new PDFRenderer(document);
			BufferedImage image;
			ByteArrayOutputStream out;
			long START, start, end, consume;
			String filename;
			for (int i = 0; i < pages; i++) {
				START = start = System.currentTimeMillis();
				image = renderer.renderImageWithDPI(i, dpi);
				out = new ByteArrayOutputStream();
				ImageIOUtil.writeImage(image, imageType.getExtensions()[0], out);
				image.getGraphics().dispose();
				bytes = out.toByteArray();
				end = System.currentTimeMillis();
				consume = end - start;
				if (maxSize > 0) {
					start = System.currentTimeMillis();
					bytes = ImageUtil.autoSize(bytes, maxSize);
					end = System.currentTimeMillis();
				}
				filename = dir + File.separator + "page_" + prependZero(i + 1, length) + "." + imageType.getExtensions()[0];
				FileUtils.writeByteArrayToFile(new File(filename), bytes);
				log.debug("\n-->[" + id + "-" + pages + "-" + (i + 1) + "] PDF 转 " + imageType.getExtensions()[0].toUpperCase()
						+ "\n   转换图片耗时 " + consume + " 毫秒"
						+ (maxSize > 0 ? "\n   缩放耗时 " + (end - start) + " 毫秒" : "")
						+ "\n   总耗时 " + (end - START) + "毫秒");
				images.add(filename);
			}
		} finally {
			if (document != null) {
				try {
					document.close();
				} catch (IOException e) {
					// do nothing
				}
			}
		}
		return images;
	}
	
	/**
	 * 前补 0 。
	 * @param num 数字。
	 * @param length 总长度。
	 * @return 补 0 后的字符串。
	 */
	private static String prependZero(int num, int length) {
		String str = String.valueOf(num);
		int diff = length - str.length();
		for (int i = 0 ; i < diff; i++) {
			str = "0" + str;
		}
		return str;
	}

}
